package com.yupi.yuaiagent.rag;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.KeywordMetadataEnricher;
import org.springframework.ai.transformer.SummaryMetadataEnricher;

import java.util.List;

/**
 * 基于 AI 的文档元信息增强器（为文档补充元信息）
 */
//@Component
public class MyKeywordEnricher {

    private final ChatModel catModel;

    public MyKeywordEnricher(ChatModel dashScopeModel) {
        this.catModel = dashScopeModel;
    }

    //关键词元信息增强器
    List<Document> enrichDocumentsByKeywords(List<Document> documents) {
        KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(this.catModel, 5);
        return keywordMetadataEnricher.apply(documents);
    }

    //摘要元信息增强器
    List<Document> enrichDocumentsBySummary(List<Document> documents) {
        return  new SummaryMetadataEnricher(catModel, List.of(SummaryMetadataEnricher.SummaryType.PREVIOUS,
                SummaryMetadataEnricher.SummaryType.CURRENT,
                SummaryMetadataEnricher.SummaryType.NEXT)).apply(documents);
    }
}
